import * as tf from '@tensorflow/tfjs'

class TensorflowService {
  constructor() {
    this.model = null;
    this.preprocessedData = null;
    this.imageSize = 224;
    this.status = '';
  }

  async loadModel(numClasses, modelType = 'mobilenet') {
    try {
      switch (modelType) {
        case 'mobilenet': {
          // 加载预训练的 MobileNet 模型
          this.status = '正在加载 MobileNet 预训练模型...';
          const mobilenet = await tf.loadLayersModel(
            'https://storage.googleapis.com/tfjs-models/tfjs/mobilenet_v1_0.25_224/model.json'
          );
          
          // 获取中间层输出
          const layer = mobilenet.getLayer('conv_pw_13_relu');
          const baseModel = tf.model({
            inputs: mobilenet.inputs,
            outputs: layer.output
          });
          
          // 冻结预训练层
          for (const layer of baseModel.layers) {
            layer.trainable = false;
          }
          
          // 构建新的分类头
          const input = tf.input({shape: [224, 224, 3]});
          const features = baseModel.apply(input);
          const globalAverage = tf.layers.globalAveragePooling2d().apply(features);
          const dropout1 = tf.layers.dropout({rate: 0.2}).apply(globalAverage);
          const dense1 = tf.layers.dense({units: 100, activation: 'relu'}).apply(dropout1);
          const dropout2 = tf.layers.dropout({rate: 0.2}).apply(dense1);
          const outputs = tf.layers.dense({units: numClasses, activation: 'softmax'}).apply(dropout2);
          
          // 创建并编译新模型
          this.model = tf.model({inputs: input, outputs: outputs});
          this.model.compile({
            optimizer: tf.train.adam(0.0001),
            loss: 'categoricalCrossentropy',
            metrics: ['accuracy']
          });
          
          this.status = 'MobileNet 模型加载完成';
          break;
        }
        
        default: {
          // 创建一个简单的 CNN 模型作为后备选项
          this.status = '创建简单 CNN 模型...';
          this.model = tf.sequential();
          
          // 添加卷积层和池化层
          this.model.add(tf.layers.conv2d({
            inputShape: [this.imageSize, this.imageSize, 3],
            filters: 32,
            kernelSize: 3,
            activation: 'relu',
            padding: 'same'
          }));
          this.model.add(tf.layers.maxPooling2d({poolSize: 2}));
          
          this.model.add(tf.layers.conv2d({
            filters: 64,
            kernelSize: 3,
            activation: 'relu',
            padding: 'same'
          }));
          this.model.add(tf.layers.maxPooling2d({poolSize: 2}));
          
          this.model.add(tf.layers.conv2d({
            filters: 64,
            kernelSize: 3,
            activation: 'relu',
            padding: 'same'
          }));
          this.model.add(tf.layers.maxPooling2d({poolSize: 2}));
          
          // 添加全连接层
          this.model.add(tf.layers.flatten());
          this.model.add(tf.layers.dense({units: 128, activation: 'relu'}));
          this.model.add(tf.layers.dropout({rate: 0.5}));
          this.model.add(tf.layers.dense({units: numClasses, activation: 'softmax'}));
          
          // 编译模型
          this.model.compile({
            optimizer: tf.train.adam(0.001),
            loss: 'categoricalCrossentropy',
            metrics: ['accuracy']
          });
          
          this.status = '简单 CNN 模型创建完成';
        }
      }
      
      return this.model;
    } catch (error) {
      console.error('加载模型时出错:', error);
      throw new Error(`加载模型失败: ${error.message}`);
    }
  }

  async preprocessImages(data) {
    const { images, numClasses, categoryMap } = data;
    
    try {
      // 创建张量数组来存储处理后的图像数据
      const processedImages = [];
      const labels = [];
      
      // 处理每张图片
      for (const image of images) {
        const img = new Image();
        img.crossOrigin = 'anonymous';  // 处理跨域问题
        img.src = image.url;
        
        await new Promise((resolve, reject) => {
          img.onload = () => {
            try {
              // 将图像转换为张量并进行预处理
              const tensor = tf.tidy(() => {
                const imageTensor = tf.browser.fromPixels(img)
                  .resizeBilinear([this.imageSize, this.imageSize])
                  .toFloat()
                  .div(255.0)
                  .expandDims();
                return imageTensor;
              });
              
              processedImages.push(tensor);
              
              // 创建one-hot编码的标签
              const labelIndex = categoryMap[image.label];
              const label = tf.oneHot(labelIndex, numClasses);
              labels.push(label);
              
              resolve();
            } catch (error) {
              reject(error);
            }
          };
          img.onerror = () => reject(new Error(`Failed to load image: ${image.name}`));
        });
      }
      
      // 合并所有图像张量和标签
      this.preprocessedData = {
        xs: tf.concat(processedImages, 0),
        ys: tf.stack(labels)
      };
      
      // 清理临时张量
      processedImages.forEach(tensor => tensor.dispose());
      labels.forEach(tensor => tensor.dispose());
      
      return this.preprocessedData;
    } catch (error) {
      console.error('预处理图像时出错:', error);
      throw error;
    }
  }

  async train(config, callbacks) {
    if (!this.model || !this.preprocessedData) {
      throw new Error('模型或数据未准备好');
    }
    
    try {
      // 配置模型
      this.model.compile({
        optimizer: tf.train.adam(config.learningRate),
        loss: 'categoricalCrossentropy',
        metrics: ['accuracy'],
      });
      
      // 开始训练
      const result = await this.model.fit(this.preprocessedData.xs, this.preprocessedData.ys, {
        epochs: config.epochs,
        batchSize: config.batchSize,
        validationSplit: config.validationSplit,
        shuffle: true,
        callbacks: {
          onEpochEnd: async (epoch, logs) => {
            if (callbacks?.onEpochEnd) {
              callbacks.onEpochEnd(epoch, logs);
            }
          }
        }
      });
      
      return result;
    } catch (error) {
      console.error('训练模型时出错:', error);
      throw error;
    }
  }

  async saveModel(modelName = 'model') {
    if (!this.model) {
      throw new Error('没有可保存的模型');
    }
    
    try {
      // 保存模型到 IndexedDB
      await this.model.save(`indexeddb://${modelName}`);
      
      // 同时提供下载选项
      await this.model.save(`downloads://${modelName}`);
      
      return true;
    } catch (error) {
      console.error('保存模型时出错:', error);
      throw error;
    }
  }

  async loadModelFromIndexedDB(modelName) {
    try {
      // 从 IndexedDB 加载模型
      this.model = await tf.loadLayersModel(`indexeddb://${modelName}`);
      return true;
    } catch (error) {
      console.error('加载模型时出错:', error);
      throw error;
    }
  }

  async listSavedModels() {
    try {
      // 获取 IndexedDB 中保存的所有模型
      const models = await tf.io.listModels();
      return Object.keys(models).map(key => ({
        name: key.replace('indexeddb://', ''),
        date: new Date(models[key].dateSaved).toLocaleString()
      }));
    } catch (error) {
      console.error('获取已保存模型列表时出错:', error);
      throw error;
    }
  }

  async deleteModel(modelName) {
    try {
      await tf.io.removeModel(`indexeddb://${modelName}`);
      return true;
    } catch (error) {
      console.error('删除模型时出错:', error);
      throw error;
    }
  }

  dispose() {
    tf.tidy(() => {
      if (this.preprocessedData) {
        this.preprocessedData.xs.dispose();
        this.preprocessedData.ys.dispose();
        this.preprocessedData = null;
      }
      if (this.model) {
        this.model.dispose();
        this.model = null;
      }
    });
  }
}

export default new TensorflowService(); 